import getopt
import sys
import oss2
import pymysql
import redis
import time
import json
import datetime
from PIL import Image
import io, os
import requests
import traceback
import base64
import random
from urllib.parse import urlparse

from aliyunsdkcore import client
from aliyunsdkgreen.request.v20180509 import ImageSyncScanRequest
import uuid

env = "prod"

opts, args = getopt.getopt(sys.argv[1:], "", "env=")
for opt, arg in opts:
    if opt == '--env':
        env = arg
        break

if env == "prod":
    mysql_host = "rm-j6ctyd19275tjf7x1.mysql.rds.aliyuncs.com"
    mysql_user = "root"
    mysql_pwd = "qwemnb0011@"
    mysql_db = "aipainter"

    redis_host = "r-j6cmt5ayr2nq1m6zpp.redis.rds.aliyuncs.com"
    redis_db = 0
    redis_pwd = "r-j6cmt5ayr2nq1m6zpp:qpalmz1209@L"

    oss_endpoint = "oss-cn-hongkong-internal.aliyuncs.com"
else:
    mysql_host = "192.168.1.56"
    mysql_user = "root"
    mysql_pwd = "ubmy8920fe101f0aabffeeaecf1d9e04"
    mysql_db = "dl_oversea"

    redis_host = "192.168.1.56"
    redis_db = 0
    redis_pwd = "ubd83922c944c7e45b45ab5c86f98e7f"

    oss_endpoint = "oss-cn-hongkong.aliyuncs.com"

# mysql
conn = pymysql.connect(host=mysql_host, user=mysql_user, password=mysql_pwd, database=mysql_db)
cur = conn.cursor()

# redis
r = redis.StrictRedis(host=redis_host, port=6379, db=redis_db, password=redis_pwd)

# oss
ossAuth = oss2.Auth('LTAI5tMzkcDHweKw7762CzQn', '6cfufuTuoQpdwR7mckhxITwkLBWtcH')
ossBucket = oss2.Bucket(ossAuth, oss_endpoint, 'aipainter-overseas')

# 图片审核
oOssValidate = client.AcsClient("LTAI5tMzkcDHweKw7762CzQn", "6cfufuTuoQpdwR7mckhxITwkLBWtcH", "cn-beijing")


# 图片审核
def ossValidate(sImgUrl):
    Request = ImageSyncScanRequest.ImageSyncScanRequest()
    Request.set_accept_format('JSON')
    aTask = {
        "dataId": str(uuid.uuid1()),
        "url": sImgUrl
    }

    Request.set_content(json.dumps({"tasks": [aTask], "scenes": ["porn"]}))
    response = oOssValidate.do_action_with_exception(Request)

    Result = json.loads(response)
    if 200 == Result["code"]:
        taskResults = Result["data"]
        for taskResult in taskResults:
            if 200 == taskResult["code"]:
                sceneResults = taskResult["results"]
                for sceneResult in sceneResults:
                    # sScene = sceneResult["scene"]
                    sSuggestion = sceneResult["suggestion"]

                    if 'pass' != sSuggestion:
                        return False
        return True
    return True


# 写入日志
def writeLog(data, machine):
    now = datetime.datetime.now().strftime('%Y-%m-%d')
    path = "../logs/{}_{}.txt".format(machine, now)
    now = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S')
    file = open(path, 'a+')
    file.write(now + " " + data + '\n')
    file.close()


def getFileBase64(imgUrl):
    imgUrl += "!webp"
    response = requests.get(imgUrl)
    base64_data = base64.b64encode(response.content).decode('utf-8')
    return base64_data


def doTask(URL, machine):
    queueKey = "painter:{}:queue".format(machine)

    loopNum = 0

    writeLog("loop start!!!", machine)

    while loopNum < 500:
        loopNum += 1

        sdinput = r.lpop(queueKey)
        if sdinput == None:
            writeLog("no task!!", machine)
            cur.execute('select 1')
            time.sleep(1)
            continue
        toStr = str(sdinput, encoding="utf-8")
        writeLog(toStr, machine)
        arr = json.loads(toStr)

        task_id = arr['task_id']
        writeLog("task_id : {} start!".format(task_id), machine)

        taskKey = "painter:task:{}".format(task_id)
        curUserId = arr['user_id'] if arr['user_id'] > 0 else arr["device_id"]
        userLenKey = "painter:user:len:{}".format(curUserId)
        pngFile = "{}.png".format(task_id)
        webpFile = "{}.webp".format(task_id)
        headers = {"User-Agent": arr["model"]}
        # drawing_image = arr['drawing_url']
        isHigh = arr['isHigh']
        style_id = arr["style_id"]

        # 默认值drawing_image
        if "drawing_url" not in arr:
            drawing_image = ""
        else:
            drawing_image = arr['drawing_url']

        # 是否是延时任务
        delay_task = False
        if "is_dealy" in arr and arr["is_delay"] == 1:
            delay_task = True

        # 设置处理中状态
        r.hset(taskKey, 'status', 1)
        try:
            if arr['type'] == 1 or arr['type'] == 2:
                writeLog("任务类型：txt2img", machine)
                sampler_index = "Euler a" if arr["type"] == 1 else "DPM++ 2M Karras"
                payload = {
                    "prompt": arr["translator_prompt"],
                    "steps": 20,
                    "cfg_scale": "7.0",
                    "width": arr["width"],
                    "height": arr["height"],
                    "negative_prompt": arr["negative_prompt"],
                    "sampler_index": sampler_index,
                }
                if isHigh == 1:
                    hig_payload = {
                        "enable_hr": True,
                        "denoising_strength": 0.5,
                        "hr_scale": 2,
                        "hr_upscaler": "Latent"
                    }
                    payload.update(hig_payload)

                payload_json = json.dumps(payload)

                response = requests.post(url=f'{URL}/sdapi/v1/txt2img', data=payload_json, headers=headers,
                                         timeout=50).json()
                image = Image.open(io.BytesIO(base64.b64decode(response['images'][0])))
                image.save(pngFile)

            else:
                writeLog("任务类型：img2img", machine)

                base64Data = getFileBase64(drawing_image)

                rand_int = random.randint(0, 1)
                interrogateModel = ""
                if rand_int == 0:
                    interrogateModel = "clip"
                else:
                    interrogateModel = "deepdanbooru"

                writeLog("interrogateModel：{}".format(interrogateModel), machine)

                payload = {
                    "image": base64Data,
                    "model": interrogateModel
                    # deepdanbooru,clip
                }

                payload_json_pro = json.dumps(payload)
                # writeLog("payload_json_pro：{}，header:{}".format(payload_json_pro,headers),machine)
                responsePro = requests.post(url=f'{URL}/sdapi/v1/interrogate', data=payload_json_pro, headers=headers,
                                            timeout=50).json()

                if "caption" not in responsePro:
                    responsePro['caption'] = ""
                arr["prompt"] = responsePro['caption']
                writeLog("图片解析的描述词：{}".format(responsePro['caption']), machine)

                base64dataArray = []
                base64dataArray.append(base64Data)

                payload = {
                    "init_images": base64dataArray,
                    "prompt": responsePro['caption'],
                    "steps": 20,
                    "denoising_strength": 0.65,
                    "width": arr["width"],
                    "height": arr["height"],
                    "negative_prompt": arr["negative_prompt"],
                    "sampler_index": "DPM++ 2M Karras",
                    "restore_faces": True
                }

                payload_json = json.dumps(payload)
                response = requests.post(url=f'{URL}/sdapi/v1/img2img', data=payload_json, headers=headers,
                                         timeout=50).json()
                image = Image.open(io.BytesIO(base64.b64decode(response['images'][0])))
                image.save(pngFile)

            # png转webp Start
            # sCmd = "ffmpeg -nostdin -i " + pngFile + "  " + webpFile
            sCmd = "ffmpeg -nostdin -i " + pngFile + " -c:v libwebp -quality 90 -compression_level 6 " + webpFile
            writeLog("图片转换: {}".format(sCmd), machine)
            os.system(sCmd)
            writeLog("png转webp成功", machine)
            # png转webp End

            # 上传图片&删除本地图片
            path = "sd/opus/%s/%s" % (time.strftime("%Y-%m-%d", time.localtime()), webpFile)
            writeLog("本地图片地址：{}".format(webpFile), machine)
            writeLog("oss图片地址：{}".format(path), machine)
            ossBucket.put_object_from_file(path, webpFile)

            # 写入数据&更新状态
            url = 'https://overseas-files.dalleai.art/' + path

            # iOS做图片审核  试用用户，如果鉴黄为黄图，就随机取出词条里的一张图给返回，正式用户的话需要返还点数
            if('1' == arr["platform"]):
                # sDefaultImgUrl = 'https://overseas-files.dalleai.art/default-audit-failure-2.png'
                bValidate = ossValidate(url)
                writeLog("OSS验证结果:{}".format(bValidate), machine)
                if(True != bValidate):

                    if ('1' == arr['is_try']):
                        url = 'https://overseas-files.dalleai.art/default-audit-failure-2.png'
                        sql = "select id from ap_model_style_special where model_id = %s"
                        cur.execute(sql,(arr['model_id']))
                        styleListSpecial = cur.fetchall()
                        if len(styleListSpecial) == 0:
                            url = 'https://overseas-files.dalleai.art/default-audit-failure-2.png'
                            
                        else:
                            targetId = random.choice(styleListSpecial)
                            sql = "select img from ap_model_style_special where id = %s"
                            cur.execute(sql,(targetId))
                            path = cur.fetchone()[0]
                            url = 'https://overseas-files.dalleai.art/' + path
                    else:
                        url = 'https://overseas-files.dalleai.art/default-audit-failure-2.png'
                        if arr['dot'] > 0:
                            sql = "update ap_user set dot=dot+%s where id = %s"
                            cur.execute(sql,(arr['dot'], arr['user_id']))
                        

                    # 删除OSS
                    writeLog("删除审核未通过图： " + path, machine)
                    ossBucket.delete_object(path)
            now_time = datetime.datetime.now()
            now_time_str = (now_time+datetime.timedelta(hours=-4)).strftime('%Y-%m-%d %H:%M:%S')
            sql = """
            insert into ap_opus(user_id, prompt, prompt_use, task_id, url, height, width, negative_prompt, type, drawing_id, drawing_image, is_new,device_id, is_high, dot, style_id, is_style_orginal, platform, is_try, machine_id, created_at)
            value(%s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s, %s)
            """
            # writeLog(sql % (arr['user_id'], arr['prompt'], arr['task_id'], url, arr['height'], arr['width'], arr['prompt'], arr['negative_prompt'], arr['type'],arr['model_id'],drawing_image,1,arr["device_id"]),machine)
            cur.execute(sql, (
            arr['user_id'], arr['prompt'], arr['translator_prompt'], arr['task_id'], url, arr['height'], arr['width'],
            arr['negative_prompt'], arr['type'], arr['model_id'], drawing_image, 1, arr["device_id"], isHigh,
            arr["dot"], style_id, arr["is_style_orginal"], arr["platform"], arr["is_try"], machine, now_time_str))

            sql = "update ap_model_list set use_num =  use_num + 1 where id = %s"
            cur.execute(sql, (arr['model_id']))

            if style_id > 0:
                sql = "update ap_model_style set use_num =  use_num + 1 where id = %s"
                cur.execute(sql, (style_id))

            conn.commit()

            # if arr['type'] == 3 :
            #    upload_obj = urlparse(drawing_image)
            #    upload_path = upload_obj.path[1:]
            #    ossBucket.delete_object(upload_path)
            #    writeLog("删除用户上传图片：{}".format(upload_path),machine)

            if not delay_task:
                r.hset(taskKey, 'status', 2)
                r.hset(taskKey, 'url', url)
                r.hset(taskKey, 'prompt', arr['prompt'])
            r.decrby(userLenKey, 1)
            r.expire(userLenKey, 600)
            writeLog("task_id : {} complate!".format(task_id), machine)

        except Exception as e:
            if not delay_task:
                r.hset(taskKey, 'status', 0)
            # r.rpush(queueKey, sdinput)
            writeLog(traceback.format_exc(), machine)

            # fail_amount 失败次数
            if ("fail_amount" not in arr):
                arr["fail_amount"] = 1
                t_sdinput = json.dumps(arr)
                r.rpush(queueKey, t_sdinput)
            else:
                url = 'https://overseas-files.dalleai.art/default.png'
                if int(arr['dot']) > 0:
                    sql = "update ap_user set dot = dot + %s where id = %s"
                    cur.execute(sql, (arr['dot'], arr['user_id']))

                    sql = "insert ap_dot_log(user_id, value, type) value(%s, %s, %s)"
                    cur.execute(sql, (arr['user_id'], arr['dot'], 5))

                    url = 'https://overseas-files.dalleai.art/default-dot.png'

                r.decrby(userLenKey, 1)
                r.expire(userLenKey, 600)

                writeLog("task : {} fail!".format(sdinput), machine)


        finally:
            if os.path.exists(pngFile):
                os.remove(pngFile)
            if os.path.exists(webpFile):
                os.remove(webpFile)

    writeLog("loop stop! have a rest!!!", machine)
